#!/usr/bin/env python3
import os
import json
from pathlib import Path
from textwrap import dedent
from typing import List, Optional
from openai import OpenAI
from pydantic import BaseModel
from dotenv import load_dotenv
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
import sys
import shutil
import logging
from datetime import datetime

# 在文件开头添加版本和作者信息
__version__ = "1.0.1"
__author__ = "Da"

# 初始化Rich控制台
console = Console()

# --------------------------------------------------------------------------------
# 1. 配置OpenAI客户端并加载环境变量
# --------------------------------------------------------------------------------
load_dotenv()  # 从.env文件加载环境变量
client = OpenAI(
    api_key=os.getenv("MY_API_KEY"),
    base_url=os.getenv("MY_BASE_URL"),
)  # 配置API_KEY和BASE_URL

# --------------------------------------------------------------------------------
# 2. 使用Pydantic定义schema以确保类型安全
# --------------------------------------------------------------------------------


# 文件创建结构
class FileToCreate(BaseModel):
    # 文件路径
    path: str
    # 文件内容
    content: str


# 文件编辑结构
class FileToEdit(BaseModel):
    # 文件路径
    path: str
    # 原始代码片段
    original_snippet: str
    # 新代码片段
    new_snippet: str


# 助手响应结构
class AssistantResponse(BaseModel):
    # 助手回复
    assistant_reply: str
    # 文件创建
    files_to_create: Optional[List[FileToCreate]] = None
    # 文件编辑
    files_to_edit: Optional[List[FileToEdit]] = None


# --------------------------------------------------------------------------------
# 3. 系统提示词
# --------------------------------------------------------------------------------
system_PROMPT = dedent("""\
    You are an elite software engineer called My Engineer with decades of experience across all programming domains.
    Your expertise spans system design, algorithms, testing, and best practices.
    You provide thoughtful, well-structured solutions while explaining your reasoning.

    Core capabilities:
    1. Code Analysis & Discussion
       - Analyze code with expert-level insight
       - Explain complex concepts clearly
       - Suggest optimizations and best practices
       - Debug issues with precision

    2. File Operations:
       a) Read existing files
          - Access user-provided file contents for context
          - Analyze multiple files to understand project structure
       
       b) Create new files
          - Generate complete new files with proper structure
          - Create complementary files (tests, configs, etc.)
       
       c) Edit existing files
          - Make precise changes using diff-based editing
          - Modify specific sections while preserving context
          - Suggest refactoring improvements

    Output Format:
    You must provide responses in this JSON structure:
    {
      "assistant_reply": "Your main explanation or response",
      "files_to_create": [
        {
          "path": "path/to/new/file",
          "content": "complete file content"
        }
      ],
      "files_to_edit": [
        {
          "path": "path/to/existing/file",
          "original_snippet": "exact code to be replaced",
          "new_snippet": "new code to insert"
        }
      ]
    }

    Guidelines:
    1. For normal responses, use 'assistant_reply'
    2. When creating files, include full content in 'files_to_create'
    3. For editing files:
       - Use 'files_to_edit' for precise changes
       - Include enough context in original_snippet to locate the change
       - Ensure new_snippet maintains proper indentation
       - Prefer targeted edits over full file replacements
    4. Always explain your changes and reasoning
    5. Consider edge cases and potential impacts
    6. Follow language-specific best practices
    7. Suggest tests or validation steps when appropriate

    Remember: You're a senior engineer - be thorough, precise, and thoughtful in your solutions.
""")

# --------------------------------------------------------------------------------
# 4. 辅助函数
# --------------------------------------------------------------------------------


# 读取本地文件
def read_local_file(file_path: str) -> str:
    """返回本地文件的文本内容"""
    with open(file_path, "r", encoding="utf-8") as f:
        return f.read()


# 创建文件
def create_file(path: str, content: str):
    """创建（或覆盖）一个文件，路径为 path 内容为 content"""
    file_path = Path(path)
    # 确保任何目录存在
    file_path.parent.mkdir(parents=True, exist_ok=True)
    with open(file_path, "w", encoding="utf-8") as f:
        f.write(content)
    # 打印创建/更新文件的消息
    console.print(f"[green]✓[/green] 创建/更新文件 '[cyan]{file_path}[/cyan]'")
    # 记录操作
    conversation_history.append(
        {"role": "assistant", "content": f"✓ Created/updated file at '{file_path}'"}
    )

    # 将实际内容添加到对话上下文
    normalized_path = normalize_path(str(file_path))
    # 记录文件内容
    conversation_history.append(
        {
            "role": "system",
            "content": f"Content of file '{normalized_path}':\n\n{content}",
        }
    )


# 显示一个包含提议的编辑并确认的表
def show_diff_table(files_to_edit: List[FileToEdit]) -> None:
    # 如果没有提议的编辑，则返回
    if not files_to_edit:
        return
    # 启用多行，通过设置show_lines=True
    table = Table(
        title="变更信息",
        show_header=True,
        header_style="bold magenta",
        show_lines=True,
    )
    # 添加列
    table.add_column("文件路径", style="cyan")
    table.add_column("原始内容", style="red")
    table.add_column("新内容", style="green")
    # 添加行
    for edit in files_to_edit:
        table.add_row(edit.path, edit.original_snippet, edit.new_snippet)
    # 打印表
    console.print(table)


# 应用diff编辑
def apply_diff_edit(path: str, original_snippet: str, new_snippet: str):
    """读取文件，替换第一个出现的original_snippet，然后覆盖"""
    try:
        content = read_local_file(path)
        if original_snippet in content:
            # 创建备份
            backup_path = backup_file(path)
            if backup_path:
                console.print(f"[green]✓[/green] 创建备份文件: '[cyan]{backup_path}[/cyan]'")
            
            updated_content = content.replace(original_snippet, new_snippet, 1)
            create_file(path, updated_content)
            
            # 添加撤销提示
            if backup_path:
                console.print(
                    "\n[yellow]提示: 如需撤销更改，可以使用备份文件恢复[/yellow]"
                )
        else:
            # 使用difflib来显示更详细的差异
            import difflib
            diff = list(difflib.unified_diff(
                content.splitlines(keepends=True),
                new_snippet.splitlines(keepends=True),
                fromfile='原始文件',
                tofile='期望更改'
            ))
            console.print("\n[yellow]详细的差异对比:[/yellow]")
            for line in diff:
                if line.startswith('+'):
                    console.print(line.rstrip(), style="green")
                elif line.startswith('-'):
                    console.print(line.rstrip(), style="red")
                else:
                    console.print(line.rstrip())
    except FileNotFoundError:
        console.print(
            f"[red]✗[/red] 文件未找到: '[cyan]{path}[/cyan]'",
            style="red",
        )


def try_handle_add_command(args: str) -> bool:
    """
    处理/add命令
    参数: 文件路径
    返回: 是否成功
    """
    if not args:
        console.print("[yellow]请指定文件路径[/yellow]")
        return True
        
    file_path = args.strip()
    try:
        # 读取文件
        content = read_local_file(file_path)
        # 将文件内容添加到对话上下文
        conversation_history.append(
            {
                "role": "system",
                "content": f"Content of file '{file_path}':\n\n{content}",
            }
        )
        console.print(
            f"[green]✓[/green] 新增文件 '[cyan]{file_path}[/cyan]' 到对话\n"
        )
        logging.info(f"添加文件到对话: {file_path}")
        return True
    except OSError as e:
        console.print(
            f"[red]✗[/red] 新增文件失败 '[cyan]{file_path}[/cyan]': {e}\n",
            style="red",
        )
        logging.error(f"添加文件失败: {file_path} - {e}")
        return True


def ensure_file_in_context(file_path: str) -> bool:
    """
    确保文件在对话上下文中
    返回True如果成功，否则返回False
    """
    try:
        # 规范化路径
        normalized_path = normalize_path(file_path)
        
        # 检查文件是否已在工作目录中
        file_marker = f"Content of file '{normalized_path}'"
        for msg in conversation_history:
            if msg["role"] == "system" and file_marker in msg["content"]:
                return True
        
        # 检查文件是否存在
        if not os.path.exists(normalized_path):
            console.print(
                f"[red]✗[/red] 文件不存在: '[cyan]{file_path}[/cyan]'",
                style="red",
            )
            return False
        
        # 读取文件
        content = read_local_file(normalized_path)
        
        # 添加到对话上下文
        conversation_history.append(
            {"role": "system", "content": f"{file_marker}:\n\n{content}"}
        )
        logging.info(f"添加文件到工作目录: {normalized_path}")
        return True
        
    except OSError as e:
        console.print(
            f"[red]✗[/red] 读取文件失败: '[cyan]{file_path}[/cyan]': {str(e)}",
            style="red",
        )
        return False


def normalize_path(path_str: str) -> str:
    """返回一个规范的绝对路径"""
    return str(Path(path_str).resolve())


# --------------------------------------------------------------------------------
# 5. 对话状态
# --------------------------------------------------------------------------------
conversation_history = [{"role": "system", "content": system_PROMPT}]

# --------------------------------------------------------------------------------
# 6. OpenAI API 交互
# --------------------------------------------------------------------------------


def guess_files_in_message(user_message: str) -> List[str]:
    """
    猜测用户消息中的文件路径
    """
    # 识别的扩展名
    recognized_extensions = [".css", ".html", ".js", ".py", ".json", ".md"]
    # 潜在路径
    potential_paths = []
    
    # 从会话历史中获取已知文件
    known_files = [
        msg["content"][msg["content"].find("'")+1:msg["content"].find("'", msg["content"].find("'")+1)]
        for msg in conversation_history 
        if msg["role"] == "system" and "Content of file '" in msg["content"]
    ]
    
    # 遍历用户消息中的每个单词
    for word in user_message.split():
        # 去除引号和标点
        path = word.strip("',\"。，；：")
        
        # 检查是否是已知文件
        if path in known_files:
            potential_paths.append(path)
            continue
            
        # 如果单词包含识别的扩展名或包含/，则认为是路径
        if any(ext in path for ext in recognized_extensions) or "/" in path:
            try:
                # 尝试规范化路径
                normalized_path = normalize_path(path)
                if os.path.exists(normalized_path):
                    potential_paths.append(normalized_path)
            except (OSError, ValueError):
                continue
    
    return potential_paths


def stream_openai_response(user_message: str):
    """流式处理大模型的响应并处理结构化输出"""
    # 猜测用户引用的文件
    potential_paths = guess_files_in_message(user_message)
    valid_files = {}

    # 尝试在API调用之前读取所有潜在的文件
    for path in potential_paths:
        try:
            content = read_local_file(path)
            valid_files[path] = content
            file_marker = f"Content of file '{path}'"
            if not any(file_marker in msg["content"] for msg in conversation_history):
                conversation_history.append(
                    {"role": "system", "content": f"{file_marker}:\n\n{content}"}
                )
        except OSError:
            error_msg = f"无法继续: 文件 '{path}' 不存在或无法访问"
            console.print(f"[red]✗[/red] {error_msg}", style="red")
            continue

    # 添加用户消息
    conversation_history.append({"role": "user", "content": user_message})
    
    try:
        # 创建流式响应
        stream = client.chat.completions.create(
            model=os.getenv("MY_MODEL"),
            messages=conversation_history,
            response_format={"type": "json_object"},
            max_completion_tokens=8000,
            stream=True,
        )

        # 打印助手回复前缀
        console.print("\nAssistant> ", style="bold blue", end="")
        
        # 完整内容
        full_content = ""
        
        # 使用spinner显示思考状态
        with console.status("[bold blue]思考中...", spinner="dots") as status:
            try:
                # 获取第一个响应
                first_chunk = next(stream)
                if first_chunk.choices[0].delta.content:
                    full_content += first_chunk.choices[0].delta.content
                    console.print(first_chunk.choices[0].delta.content, end="")
                
                # 关闭状态显示
                status.stop()
                
                # 继续处理剩余的响应
                for chunk in stream:
                    if chunk.choices[0].delta.content:
                        content_chunk = chunk.choices[0].delta.content
                        full_content += content_chunk
                        console.print(content_chunk, end="", highlight=False)
                        
            except StopIteration:
                # 处理流结束的情况
                pass
            except Exception as e:
                console.print(f"\n[red]✗[/red] 处理响应时出错: {str(e)}")
                return AssistantResponse(assistant_reply="处理响应时出错", files_to_create=[])
        
        console.print()  # 换行

        try:
            # 解析JSON
            parsed_response = json.loads(full_content)
            
            # 确保assistant_reply存在
            if "assistant_reply" not in parsed_response:
                parsed_response["assistant_reply"] = ""

            # 处理文件编辑...
            if "files_to_edit" in parsed_response and parsed_response["files_to_edit"]:
                new_files_to_edit = []
                for edit in parsed_response["files_to_edit"]:
                    try:
                        edit_abs_path = normalize_path(edit["path"])
                        if edit_abs_path in valid_files or ensure_file_in_context(
                            edit_abs_path
                        ):
                            edit["path"] = edit_abs_path
                            new_files_to_edit.append(edit)
                    except (OSError, ValueError):
                        console.print(
                            f"[yellow]⚠[/yellow] 跳过无效路径 '{edit['path']}'",
                            style="yellow",
                        )
                        continue
                parsed_response["files_to_edit"] = new_files_to_edit

            response_obj = AssistantResponse(**parsed_response)

            # 将助手的文本回复添加到对话上下文
            conversation_history.append(
                {"role": "assistant", "content": response_obj.assistant_reply}
            )

            return response_obj

        except json.JSONDecodeError:
            error_msg = "解析响应JSON失败"
            console.print(f"[red]✗[/red] {error_msg}", style="red")
            return AssistantResponse(assistant_reply=error_msg, files_to_create=[])

    except Exception as e:
        error_msg = f"请检查配置是否正确: {str(e)}"
        console.print(f"\n[red]✗[/red] {error_msg}", style="red")
        return AssistantResponse(assistant_reply=error_msg, files_to_create=[])


# --------------------------------------------------------------------------------
# 7. 主交互循环
# --------------------------------------------------------------------------------


def check_environment():
    """检查必要的环境变量是否存在"""
    required_vars = ["MY_API_KEY", "MY_BASE_URL", "MY_MODEL"]
    missing_vars = [var for var in required_vars if not os.getenv(var)]
    if missing_vars:
        console.print(f"[red]错误: 缺少必要的环境变量: {', '.join(missing_vars)}[/red]")
        sys.exit(1)


# 添加备份管理器类
class BackupManager:
    def __init__(self):
        self._backups = {}  # 文件路径 -> 备份路径的映射
        self._backup_time = {}  # 文件路径 -> 备份时间的映射
        
    def create_backup(self, file_path: str) -> Optional[str]:
        """创建文件备份，如果最近已备份则返回现有备份"""
        file_path = str(Path(file_path).resolve())
        current_time = datetime.now()
        
        # 检查是否已有最近的备份（5分钟内）
        if file_path in self._backup_time:
            time_diff = (current_time - self._backup_time[file_path]).total_seconds()
            if time_diff < 300:  # 5分钟内
                return self._backups[file_path]
        
        try:
            backup_path = f"{file_path}.bak"
            shutil.copy2(file_path, backup_path)
            self._backups[file_path] = backup_path
            self._backup_time[file_path] = current_time
            logging.info(f"创建备份: {file_path} -> {backup_path}")
            return backup_path
        except OSError as e:
            logging.error(f"创建备份失败 {file_path}: {e}")
            console.print(f"[yellow]警告: 无法创建文件备份: {e}[/yellow]")
            return None
    
    def get_backup(self, file_path: str) -> Optional[str]:
        """获取文件的备份路径"""
        file_path = str(Path(file_path).resolve())
        return self._backups.get(file_path)
    
    def clear_old_backups(self):
        """清理超过5分钟的备份记录"""
        current_time = datetime.now()
        old_paths = [
            path for path, time in self._backup_time.items()
            if (current_time - time).total_seconds() >= 300
        ]
        for path in old_paths:
            self._backups.pop(path, None)
            self._backup_time.pop(path, None)

# 创建全局备份管理器实例
backup_manager = BackupManager()

# 修改备份文件函数
def backup_file(file_path: str) -> Optional[str]:
    """在修改文件前创建备份"""
    return backup_manager.create_backup(file_path)

# 修改清理备份函数
def clean_backups() -> bool:
    """清理备份文件"""
    try:
        # 清理旧的备份记录
        backup_manager.clear_old_backups()
        
        # 查找所有.bak文件
        backup_files = list(Path('.').glob('*.bak'))
        if not backup_files:
            console.print("[yellow]没有找到备份文件[/yellow]")
            return True
        
        # 显示要删除的文件
        console.print("\n[bold]将删除以下备份文件:[/bold]")
        for file in backup_files:
            console.print(f"[cyan]•[/cyan] [green]{file}[/green]")
        
        # 确认删除
        confirm = console.input("\n确认删除这些文件? ([green]y[/green]/[red]n[/red]): ").strip().lower()
        if confirm != 'y':
            console.print("[yellow]取消删除[/yellow]")
            return True
        
        # 删除文件
        for file in backup_files:
            file.unlink()
            console.print(f"[green]✓[/green] 删除文件: '{file}'")
        
        # 清空备份记录
        backup_manager._backups.clear()
        backup_manager._backup_time.clear()
        
        console.print(f"\n[green]✓[/green] 已清理 {len(backup_files)} 个备份文件")
        logging.info(f"清理了 {len(backup_files)} 个备份文件")
        
    except Exception as e:
        console.print(f"[red]✗[/red] 清理备份失败: {str(e)}")
        logging.error(f"清理备份失败: {str(e)}")
    
    return True


# --------------------------------------------------------------------------------
# 1. 日志配置
# --------------------------------------------------------------------------------
def setup_logging():
    """配置日志"""
    log_file = f"my_engineer_{datetime.now().strftime('%Y%m%d')}.log"
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file, encoding='utf-8'),
            logging.StreamHandler()
        ]
    )


# --------------------------------------------------------------------------------
# 2. 改进文件操作
# --------------------------------------------------------------------------------
def safe_create_file(path: str, content: str) -> bool:
    """安全地创建文件，包含错误处理和备份"""
    try:
        file_path = Path(path)
        
        # 检查路径是否包含非法字符
        if any(c in str(file_path) for c in '<>:"|?*\\'):
            raise ValueError("文件路径包含非法字符")
        
        # 检查文件是否已存在
        if file_path.exists():
            # 创建备份
            backup_path = backup_file(str(file_path))
            if not backup_path:
                return False
        
        # 确保目录存在
        file_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 检查磁盘空间
        if shutil.disk_usage(file_path.parent).free < len(content) * 2:
            raise OSError("磁盘空间不足")
        
        # 写入文件
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(content)
            
        logging.info(f"成功创建/更新文件: {file_path}")
        console.print(f"[green]✓[/green] 创建/更新文件 '[cyan]{file_path}[/cyan]'")
        
        # 自动将文件添加到工作目录
        add_file_to_workspace(str(file_path))
        return True
        
    except ValueError as e:
        logging.error(f"创建文件失败(非法路径) {path}: {str(e)}")
        console.print("[red]✗[/red] 创建文件失败: 非法的文件路径")
    except OSError as e:
        logging.error(f"创建文件失败(IO错误) {path}: {str(e)}")
        console.print(f"[red]✗[/red] 创建文件失败: {str(e)}")
    except Exception as e:
        logging.error(f"创建文件失败(未知错误) {path}: {str(e)}")
        console.print("[red]✗[/red] 创建文件失败: 未知错误")
    return False

def add_file_to_workspace(file_path: str) -> None:
    """将文件添加到工作目录"""
    try:
        # 读取文件内容
        content = read_local_file(file_path)
        # 文件标记
        file_marker = f"Content of file '{file_path}'"
        
        # 检查文件是否已在工作目录中
        for msg in conversation_history:
            if msg["role"] == "system" and file_marker in msg["content"]:
                # 更新现有文件内容
                msg["content"] = f"{file_marker}:\n\n{content}"
                logging.info(f"更新工作目录中的文件: {file_path}")
                return
        
        # 如果文件不在工作目录中，添加它
        conversation_history.append(
            {"role": "system", "content": f"{file_marker}:\n\n{content}"}
        )
        logging.info(f"添加文件到工作目录: {file_path}")
        
    except Exception as e:
        logging.error(f"添加文件到工作目录失败 {file_path}: {str(e)}")
        console.print(f"[yellow]警告: 无法将文件添加到工作目录: {str(e)}[/yellow]")


# --------------------------------------------------------------------------------
# 3. 改进用户界面
# --------------------------------------------------------------------------------
def show_welcome():
    """显示欢迎信息和帮助"""
    version = "1.0.1"
    model = os.getenv("MY_MODEL", "未指定")
    
    # 获取所有命令的描述
    commands = get_commands()
    command_text = ""
    for cmd_name, cmd in commands.items():
        args = f" [cyan]{cmd.args_description}[/cyan]" if cmd.args_description else ""
        padding = " " * (15 - len(cmd_name) - len(args))
        command_text += f"[magenta]{cmd_name}[/magenta]{args}{padding}{cmd.description}\n"
    
    welcome_text = f"""[bold blue]My Engineer[/bold blue] v{version}
[dim]使用模型: {model}[/dim]

[bold white]支持的命令:[/bold white]
{command_text}
[bold white]特性:[/bold white]
• 自动文件备份
• 详细的差异对比
• 操作日志记录
• 多文件上下文
• 实时状态监控
"""
    console.print(Panel(welcome_text, border_style="blue"))

def handle_command(cmd: str) -> bool:
    """处理特殊命令"""
    cmd = cmd.strip().lower()
    if not cmd.startswith('/'):
        return True
        
    # 分割命令和参数
    parts = cmd.split(maxsplit=1)
    command = parts[0]
    args = parts[1] if len(parts) > 1 else ""
    
    commands = get_commands()
    if command in commands:
        result = commands[command].handler(args)
        return True if result is None else result
        
    # 未知命令
    console.print(f"[yellow]未知命令: {command}[/yellow]")
    console.print("使用 [magenta]/help[/magenta] 查看可用命令")
    return True

def clear_conversation():
    """清除对话历史和屏幕"""
    # 清除对话历史
    conversation_history.clear()
    conversation_history.append({"role": "system", "content": system_PROMPT})
    
    # 清空屏幕
    console.clear()
    
    # 显示状态
    show_status()
    
    console.print("[green]✓[/green] 对话历史已清除")
    logging.info("对话历史已清除")

def show_history():
    """显示对话历史"""
    table = Table(title="对话历史")
    table.add_column("角色", style="cyan")
    table.add_column("内容", style="white", overflow="fold")
    
    for msg in conversation_history[1:]:  # 跳过系统提示
        role = msg["role"]
        content = msg["content"]
        if len(content) > 100:
            content = content[:100] + "..."
        table.add_row(role, content)
    
    console.print(table)

def save_history(filename: str):
    """保存对话历史到文件"""
    if not filename.endswith('.json'):
        filename += '.json'
    
    try:
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(conversation_history, f, ensure_ascii=False, indent=2)
        console.print(f"[green]✓[/green] 对话历史已保存到 '{filename}'")
        logging.info(f"对话历史已保存到 {filename}")
    except Exception as e:
        console.print(f"[red]✗[/red] 保存失败: {str(e)}")
        logging.error(f"保存对话历史失败: {str(e)}")

def undo_last_change():
    """撤销最后一次文件修改"""
    # 获取所有有备份的文件及其备份时间
    backup_files = [
        (path, time) 
        for path, time in backup_manager._backup_time.items()
        if backup_manager.get_backup(path)
    ]
    
    if not backup_files:
        console.print("[yellow]没有可撤销的修改[/yellow]")
        return
    
    # 找到最近修改的文件
    latest_file, _ = max(backup_files, key=lambda x: x[1])
    backup_path = backup_manager.get_backup(latest_file)
    
    try:
        # 恢复文件
        shutil.copy2(backup_path, latest_file)
        # 删除备份文件
        Path(backup_path).unlink()
        # 从备份管理器中移除记录
        backup_manager._backups.pop(latest_file, None)
        backup_manager._backup_time.pop(latest_file, None)
        
        console.print(f"[green]✓[/green] 已恢复文件 '[cyan]{latest_file}[/cyan]'")
        logging.info(f"已撤销对 {latest_file} 的修改")
        
        # 更新工作目录中的文件内容
        add_file_to_workspace(latest_file)
        
    except Exception as e:
        console.print(f"[red]✗[/red] 撤销失败: {str(e)}")
        logging.error(f"撤销修改失败: {str(e)}")

def list_session_files():
    """列出当前会话中的文件"""
    files = get_workspace_files()
    
    if not files:
        console.print("[yellow]当前会话中没有可编辑的文件[/yellow]")
        return

    # 创建表格
    table = Table(title="当前会话可编辑的文件")
    table.add_column("序号", style="cyan", justify="right")
    table.add_column("文件路径", style="green")
    table.add_column("大小", style="magenta", justify="right")
    table.add_column("修改时间", style="blue")
    table.add_column("状态", style="yellow")

    # 添加文件信息
    for i, file_path in enumerate(files, 1):
        path = Path(file_path)
        try:
            stats = path.stat()
            size = format_size(stats.st_size)
            mtime = datetime.fromtimestamp(stats.st_mtime).strftime("%Y-%m-%d %H:%M:%S")
            has_backup = backup_manager.get_backup(file_path) is not None
            status = "有备份" if has_backup else "原始"
            table.add_row(str(i), file_path, size, mtime, status)
        except OSError:
            table.add_row(str(i), file_path, "访问失败", "未知", "错误")

    console.print(table)

def format_size(size: int) -> str:
    """格式化文件大小"""
    for unit in ['B', 'KB', 'MB', 'GB']:
        if size < 1024:
            return f"{size:.1f}{unit}"
        size /= 1024
    return f"{size:.1f}TB"

# --------------------------------------------------------------------------------
# 4. 主函数改进
# --------------------------------------------------------------------------------
def main():
    try:
        # 设置日志
        setup_logging()
        logging.info("程序启动")
        
        # 检查环境变量
        check_environment()
        
        # 显示欢迎信息
        show_welcome()
        
        while True:
            try:
                user_input = console.input("[bold green]You>[/bold green] ").strip()
            except (EOFError, KeyboardInterrupt):
                console.print("\n[yellow]正在安全退出...[/yellow]")
                handle_exit()
                break
                
            if not user_input:
                continue
                
            # 处理命令
            if not handle_command(user_input):
                break
                
            # 处理常规输入
            if not user_input.startswith('/'):
                response_data = stream_openai_response(user_input)
                handle_response(response_data)
                
    except Exception as e:
        logging.error(f"程序异常: {str(e)}", exc_info=True)
        console.print(f"[red]发生错误: {str(e)}[/red]")
        console.print_exception()
    finally:
        handle_exit()

def handle_response(response_data: AssistantResponse):
    """处理AI响应"""
    if response_data.files_to_create:
        for file_info in response_data.files_to_create:
            safe_create_file(file_info.path, file_info.content)

    if response_data.files_to_edit:
        show_diff_table(response_data.files_to_edit)
        confirm = console.input("\n应用这些更改吗? ([green]y[/green]/[red]n[/red]): ").strip().lower()
        
        if confirm == "y":
            # 按文件分组编辑操作
            edits_by_file = {}
            for edit_info in response_data.files_to_edit:
                if edit_info.path not in edits_by_file:
                    edits_by_file[edit_info.path] = []
                edits_by_file[edit_info.path].append(edit_info)
            
            # 对每个文件进行一次性修改
            for file_path, edits in edits_by_file.items():
                try:
                    # 读取文件内容
                    content = read_local_file(file_path)
                    
                    # 创建备份（每个文件只备份一次）
                    backup_path = backup_file(file_path)
                    if backup_path:
                        console.print(f"[green]✓[/green] 创建备份文件: '[cyan]{backup_path}[/cyan]'")
                    
                    # 应用所有编辑
                    modified_content = content
                    for edit in edits:
                        if edit.original_snippet in modified_content:
                            modified_content = modified_content.replace(
                                edit.original_snippet, 
                                edit.new_snippet, 
                                1
                            )
                        else:
                            console.print(
                                f"[yellow]⚠[/yellow] 在文件 '{file_path}' 中未找到要替换的代码片段"
                            )
                    
                    # 写入修改后的内容
                    with open(file_path, "w", encoding="utf-8") as f:
                        f.write(modified_content)
                    
                    console.print(f"[green]✓[/green] 更新文件: '[cyan]{file_path}[/cyan]'")
                    
                    # 更新工作目录
                    add_file_to_workspace(file_path)
                    
                except Exception as e:
                    console.print(f"[red]✗[/red] 处理文件 '{file_path}' 失败: {str(e)}")
                    logging.error(f"处理文件失败 {file_path}: {str(e)}")
            
            # 添加撤销提示
            if any(backup_manager.get_backup(path) for path in edits_by_file.keys()):
                console.print("\n[yellow]提示: 如需撤销更改，可以使用 /undo 命令[/yellow]")
        else:
            console.print("[yellow]跳过应用差异编辑[/yellow]")
            logging.info("用户取消了文件修改")


# 添加新的命令处理类
class Command:
    def __init__(self, name: str, description: str, handler, args_description: str = ""):
        self.name = name
        self.description = description
        self.handler = handler
        self.args_description = args_description

def get_commands() -> dict:
    """返回所有支持的命令"""
    return {
        # 基础命令
        "/help": Command("help", "显示帮助信息", lambda _: show_welcome()),
        "/clear": Command("clear", "清除对话历史和屏幕", lambda _: clear_conversation()),
        "/exit": Command("exit", "退出程序", lambda _: False),
        "/quit": Command("quit", "退出程序", lambda _: False),
        
        # 文件操作命令
        "/ls": Command("ls", "列出当前会话可编辑的文件", lambda _: list_session_files()),
        "/pwd": Command("pwd", "显示当前工作目录", lambda _: show_working_directory()),
        "/add": Command("add", "导入文件到对话", try_handle_add_command, "<文件路径>"),
        "/rm": Command("rm", "从会话中移除文件", remove_session_file, "<序号或文件路径>"),
        "/cat": Command("cat", "显示文件内容", show_file_content, "<序号或文件路径>"),
        "/find": Command("find", "搜索文件内容", search_file_content, "<关键词>"),
        
        # 会话管理命令
        "/history": Command("history", "显示对话历史", lambda _: show_history()),
        "/save": Command("save", "保存对话历史", save_history, "<文件名>"),
        "/load": Command("load", "加载历史对话", load_history, "<文件名>"),
        "/undo": Command("undo", "撤销上次文件修改", lambda _: undo_last_change()),
        "/clean": Command("clean", "清理备份文件", lambda _: clean_backups()),
        
        # 系统信息命令
        "/version": Command("version", "显示版本信息", lambda _: show_version()),
        "/status": Command("status", "显示当前状态", lambda _: show_status()),
        "/info": Command("info", "显示系统信息", lambda _: show_system_info()),
    }

def show_version():
    """显示版本信息"""
    version = "1.0.1"
    console.print(f"[bold blue]My Engineer[/bold blue] v{version}")
    console.print(f"[dim]模型: {os.getenv('MY_MODEL', '未指定')}[/dim]")

def show_status():
    """显示当前状态"""
    # 计算对话轮数
    turns = len([msg for msg in conversation_history if msg["role"] == "user"])
    # 获取最后一次操作时间
    last_operation = "无"
    if len(conversation_history) > 1:
        last_operation = datetime.now().strftime("%H:%M:%S")
    
    status_text = f"""[bold white]当前状态:[/bold white]
• 对话轮数: {turns}
• 最后操作: {last_operation}
• 已加载文件: {len([msg for msg in conversation_history if msg["role"] == "system" and "Content of file" in msg["content"]])}
• 日志文件: my_engineer_{datetime.now().strftime('%Y%m%d')}.log
"""
    console.print(Panel(status_text, border_style="blue"))

def show_working_directory():
    """显示当前工作目录"""
    try:
        cwd = Path.cwd()
        # 获取父目录
        parent = cwd.parent
        
        # 创建表格
        table = Table(title="当前工作目录")
        table.add_column("类型", style="cyan")
        table.add_column("路径", style="green")
        
        # 添加当前目录和父目录信息
        table.add_row("当前目录", str(cwd))
        table.add_row("父目录", str(parent))
        
        console.print(table)
        
        # 显示目录中的文件数量
        file_count = len(list(cwd.glob('*')))
        console.print(f"\n当前目录包含 [cyan]{file_count}[/cyan] 个文件/目录")
        
    except Exception as e:
        console.print(f"[red]✗[/red] 获取工作目录失败: {str(e)}")

def remove_session_file(args: str) -> bool:
    """从会话中移除文件"""
    if not args:
        # 如果没有提供参数，显示当前可移除的文件列表
        files = []
        for msg in conversation_history:
            if msg["role"] == "system" and "Content of file '" in msg["content"]:
                start = msg["content"].find("'") + 1
                end = msg["content"].find("'", start)
                files.append(msg["content"][start:end])
        
        if not files:
            console.print("[yellow]当前会话中没有可移除的文件[/yellow]")
            return True
            
        console.print("\n[bold]可移除的文件:[/bold]")
        for i, file in enumerate(files, 1):
            console.print(f"[cyan]{i}.[/cyan] [green]{file}[/green]")
        console.print("\n使用 [magenta]/rm <序号或文件路径>[/magenta] 移除文件")
        return True
    
    try:
        # 尝试解析为序号
        if args.isdigit():
            idx = int(args) - 1
            files = [msg["content"][msg["content"].find("'")+1:msg["content"].find("'", msg["content"].find("'")+1)]
                    for msg in conversation_history 
                    if msg["role"] == "system" and "Content of file '" in msg["content"]]
            
            if 0 <= idx < len(files):
                file_to_remove = files[idx]
            else:
                console.print("[red]✗[/red] 无效的序号")
                return True
        else:
            file_to_remove = args.strip()
        
        # 从会话历史中移除文件
        new_history = []
        removed = False
        for msg in conversation_history:
            if msg["role"] == "system" and "Content of file '" in msg["content"]:
                if file_to_remove in msg["content"]:
                    removed = True
                    continue
            new_history.append(msg)
        
        if removed:
            conversation_history.clear()
            conversation_history.extend(new_history)
            console.print(f"[green]✓[/green] 已从会话中移除文件: '[cyan]{file_to_remove}[/cyan]'")
            logging.info(f"从会话中移除文件: {file_to_remove}")
        else:
            console.print(f"[yellow]未找到文件: '{file_to_remove}'[/yellow]")
        
    except Exception as e:
        console.print(f"[red]✗[/red] 移除文件失败: {str(e)}")
        logging.error(f"移除文件失败: {str(e)}")
    
    return True

def show_file_content(args: str) -> bool:
    """显示文件内容"""
    if not args:
        console.print("[yellow]请指定文件序号或路径[/yellow]")
        return True
    
    try:
        # 获取文件路径
        file_path = get_file_path(args)
        if not file_path:
            return True
            
        # 读取文件内容
        content = read_local_file(file_path)
        
        # 创建面板显示内容
        console.print(Panel(
            content,
            title=f"[blue]{file_path}[/blue]",
            border_style="blue",
            highlight=True
        ))
        
    except Exception as e:
        console.print(f"[red]✗[/red] 读取文件失败: {str(e)}")
    
    return True

def search_file_content(args: str) -> bool:
    """搜索文件内容"""
    if not args:
        console.print("[yellow]请指定搜索关键词[/yellow]")
        return True
    
    keyword = args.lower()
    results = []
    
    # 搜索所有已加载的文件
    for msg in conversation_history:
        if msg["role"] == "system" and "Content of file '" in msg["content"]:
            file_path = msg["content"][msg["content"].find("'")+1:msg["content"].find("'", msg["content"].find("'")+1)]
            content = msg["content"].split("\n\n", 1)[1]
            
            # 搜索每一行
            for i, line in enumerate(content.split('\n'), 1):
                if keyword in line.lower():
                    results.append((file_path, i, line.strip()))
    
    if not results:
        console.print(f"[yellow]未找到包含 '{keyword}' 的内容[/yellow]")
        return True
    
    # 显示搜索结果
    table = Table(title=f"搜索结果: '{keyword}'")
    table.add_column("文件", style="green")
    table.add_column("行号", style="cyan", justify="right")
    table.add_column("内容", style="white")
    
    for file_path, line_num, line in results:
        table.add_row(file_path, str(line_num), line)
    
    console.print(table)
    return True

def load_history(args: str) -> bool:
    """加载历史对话"""
    if not args:
        console.print("[yellow]请指定历史文件路径[/yellow]")
        return True
    
    file_path = args.strip()
    if not file_path.endswith('.json'):
        file_path += '.json'
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            history = json.load(f)
        
        # 验证历史记录格式
        if not isinstance(history, list) or not all(isinstance(msg, dict) for msg in history):
            raise ValueError("无效的历史记录格式")
        
        # 更新对话历史
        conversation_history.clear()
        conversation_history.extend(history)
        
        console.print(f"[green]✓[/green] 已加载历史对话: '{file_path}'")
        logging.info(f"加载历史对话: {file_path}")
        
    except Exception as e:
        console.print(f"[red]✗[/red] 加载历史失败: {str(e)}")
        logging.error(f"加载历史失败: {str(e)}")
    
    return True

def show_system_info():
    """显示系统信息"""
    import platform
    import psutil
    
    try:
        workspace_files = get_workspace_files()
        system_info = {
            "操作系统": f"{platform.system()} {platform.release()}",
            "Python版本": platform.python_version(),
            "CPU使用率": f"{psutil.cpu_percent()}%",
            "内存使用": f"{psutil.Process().memory_info().rss / 1024 / 1024:.1f}MB",
            "工作目录": str(Path.cwd()),
            "会话文件数": len(workspace_files),
            "对话轮数": len([msg for msg in conversation_history if msg["role"] == "user"]),
            "备份文件数": len(list(Path('.').glob('*.bak'))),
        }
        
        table = Table(title="系统信息")
        table.add_column("项目", style="cyan")
        table.add_column("值", style="green")
        
        for key, value in system_info.items():
            table.add_row(key, str(value))
        
        console.print(table)
        
    except Exception as e:
        console.print(f"[red]✗[/red] 获取系统信息失败: {str(e)}")

def extract_file_path_from_content(content: str) -> Optional[str]:
    """从消息内容中提取文件路径"""
    if "Content of file '" not in content:
        return None
    start = content.find("'") + 1
    end = content.find("'", start)
    return content[start:end] if start > 0 and end > start else None

def get_workspace_files() -> List[str]:
    """获取当前工作目录中的所有文件"""
    return [
        path for msg in conversation_history 
        if msg["role"] == "system" 
        and (path := extract_file_path_from_content(msg["content"])) is not None
    ]

def get_file_path(args: str) -> Optional[str]:
    """从参数中获取文件路径（支持序号或直接路径）"""
    if args.isdigit():
        idx = int(args) - 1
        files = get_workspace_files()
        if 0 <= idx < len(files):
            return files[idx]
        console.print("[red]✗[/red] 无效的序号")
        return None
    return args.strip()

# 添加优雅退出处理
def handle_exit():
    """处理程序退出"""
    # 检查是否有未保存的对话
    if len(conversation_history) > 1:  # 超过系统提示的消息
        console.print("\n[yellow]检测到未保存的对话[/yellow]")
        save = console.input("是否保存? ([green]y[/green]/[red]n[/red]): ").strip().lower()
        if save == 'y':
            filename = f"conversation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
            save_history(filename)
    
    # 检查是否有备份文件
    backup_files = list(Path('.').glob('*.bak'))
    if backup_files:
        console.print(f"\n[yellow]存在 {len(backup_files)} 个备份文件[/yellow]")
        clean = console.input("是否清理? ([green]y[/green]/[red]n[/red]): ").strip().lower()
        if clean == 'y':
            clean_backups()
    
    console.print("\n[blue]感谢使用 My Engineer![/blue]")
    logging.info("程序正常退出")


if __name__ == "__main__":
    main()
